'use client';

import { useChat, Message } from '@ai-sdk/react';
import { useState, useEffect, useRef } from 'react';
import ReactMarkdown from 'react-markdown';
import { GeistSans } from 'geist/font/sans';
import { GeistMono } from 'geist/font/mono';

export default function Home() {
  const [isThinkingExpanded, setIsThinkingExpanded] = useState(true);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const { messages, input, handleInputChange, handleSubmit } = useChat({
    api: 'http://localhost:8000/chat/stream',
    headers: {
      'Content-Type': 'application/json',
    },
    body: {
      stream: true
    }
  });

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  useEffect(() => {
    scrollToBottom();
  }, [messages]);

  const parseMessageContent = (content: string) => {
    if (content.includes('</think>')) {
      const [thinking, ...rest] = content.split('</think>');
      return {
        thinking: thinking.replace('<think>', '').trim(),
        finalResponse: rest.join('</think>').trim(),
        isComplete: true
      };
    }
    if (content.includes('<think>')) {
      return {
        thinking: content.replace('<think>', '').trim(),
        finalResponse: '',
        isComplete: false
      };
    }
    return {
      thinking: '',
      finalResponse: content,
      isComplete: true
    };
  };

  const isAnyMessageIncomplete = messages.some((message: Message) => 
    message.role === 'assistant' && !parseMessageContent(message.content).isComplete
  );

  return (
    <div className={`flex flex-col min-h-screen bg-gray-50 ${GeistSans.className}`}>
      {/* Header */}
      <header className="bg-white shadow-sm">
        <div className="max-w-4xl mx-auto px-4 py-4">
          <h1 className="text-xl font-semibold text-gray-900">DeepSeek Chat</h1>
        </div>
      </header>

      {/* Main Chat Area */}
      <main className="flex-1 max-w-4xl mx-auto w-full px-4 py-6 space-y-6">
        <div className="space-y-6">
          {messages.map((message: Message) => (
            <div key={message.id} className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}>
              <div className={`rounded-lg py-3 px-4 max-w-[85%] ${
                message.role === 'user'
                  ? 'bg-blue-600 text-white'
                  : 'bg-white text-gray-900 shadow-sm'
              }`}>
                {message.role === 'assistant' ? (
                  <>
                    {(() => {
                      const { thinking, finalResponse, isComplete } = parseMessageContent(message.content);
                      return (
                        <>
                          {(thinking || !isComplete) && (
                            <div className="mb-4 space-y-2">
                              <div className="flex items-center gap-2">
                                <button 
                                  onClick={() => setIsThinkingExpanded(!isThinkingExpanded)}
                                  className="flex items-center gap-2 text-sm text-gray-500 hover:text-gray-700"
                                >
                                  <svg 
                                    className={`w-4 h-4 transform transition-transform ${isThinkingExpanded ? 'rotate-0' : '-rotate-180'}`} 
                                    fill="none" 
                                    viewBox="0 0 24 24" 
                                    stroke="currentColor"
                                  >
                                    <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M5 15l7-7 7 7" />
                                  </svg>
                                  <span>Thinking Process</span>
                                </button>
                              </div>
                              {isThinkingExpanded && (
                                <div className="pl-4 border-l-2 border-gray-200">
                                  <div className={`text-sm text-gray-600 whitespace-pre-wrap ${GeistMono.className}`}>{thinking}</div>
                                </div>
                              )}
                            </div>
                          )}
                          {isComplete && finalResponse && (
                            <div className="prose prose-sm max-w-none">
                              <ReactMarkdown>{finalResponse}</ReactMarkdown>
                            </div>
                          )}
                        </>
                      );
                    })()}
                  </>
                ) : (
                  <div className="whitespace-pre-wrap">{message.content}</div>
                )}
              </div>
            </div>
          ))}
          <div ref={messagesEndRef} />
        </div>
      </main>

      {/* Input Area */}
      <div className="bg-white border-t">
        <div className="max-w-4xl mx-auto px-4 py-4">
          <form onSubmit={handleSubmit} className="flex gap-2">
            <input
              type="text"
              value={input}
              onChange={handleInputChange}
              placeholder="Type your message..."
              className="flex-1 rounded-lg border border-gray-300 px-4 py-2 focus:outline-none focus:ring-2 focus:ring-blue-500"
              disabled={isAnyMessageIncomplete}
            />
            <button
              type="submit"
              disabled={!input.trim() || isAnyMessageIncomplete}
              className="px-6 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:opacity-50 disabled:cursor-not-allowed"
            >
              {isAnyMessageIncomplete ? 'Sending...' : 'Send'}
            </button>
          </form>
        </div>
      </div>
    </div>
  );
}
